import os
import subprocess
import logging

class ADBScreenshot:
    ROBOT_LIBRARY_SCOPE = "GLOBAL"  # 设置库的作用域为全局

    def __init__(self, device_id=None):
        """
        初始化ADB工具。
        :param device_id: 设备ID（可选，如果连接了多个设备）
        """
        self.device_id = device_id
        self.adb_command = ["adb"]
        if device_id:
            self.adb_command.extend(["-s", device_id])
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
        logging.info("ADB Screenshot Library 初始化完成。")

    def run_adb_command(self, command):
        """
        执行ADB命令。
        :param command: ADB命令列表
        :return: 命令的输出
        :raises: 如果命令执行失败，抛出异常并记录日志
        """
        full_command = self.adb_command + command
        logging.info(f"正在执行ADB命令: {' '.join(full_command)}")
        try:
            result = subprocess.run(full_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            if result.returncode != 0:
                logging.error(f"ADB命令执行失败: {result.stderr.decode()}")
                raise Exception(f"ADB命令执行失败: {result.stderr.decode()}")
            logging.info(f"ADB命令执行成功，输出: {result.stdout.decode()}")
            return result.stdout.decode()
        except Exception as e:
            logging.error(f"运行ADB命令时发生异常: {e}")
            raise

    def take_full_screenshot(self, output_path="/sdcard/fullscreen.png"):
        """
        捕获设备的全屏截图。
        :param output_path: 截图保存路径（设备上的路径）
        :raises: 如果截图失败，抛出异常并记录日志
        """
        logging.info(f"开始捕获全屏截图，保存路径: {output_path}")
        try:
            self.run_adb_command(["shell", "screencap", "-p", output_path])
            logging.info(f"全屏截图已成功保存到设备路径: {output_path}")
        except Exception as e:
            logging.error(f"捕获全屏截图失败: {e}")
            raise

    def crop_screenshot(self, input_path, output_path, x, y, width, height):
        """
        裁剪截图。
        :param input_path: 输入文件路径（设备上的全屏截图路径）
        :param output_path: 输出文件路径（设备上的裁剪截图路径）
        :param x: 裁剪区域的左上角X坐标
        :param y: 裁剪区域的左上角Y坐标
        :param width: 裁剪区域的宽度
        :param height: 裁剪区域的高度
        :raises: 如果裁剪失败，抛出异常并记录日志
        """
        logging.info(f"开始裁剪截图，输入路径: {input_path}，输出路径: {output_path}")
        try:
            self.run_adb_command([
                "shell", "dd", "if={}".format(input_path), "of={}".format(output_path),
                "bs=1", "skip={}".format(x * 4 + y * 4 * width), "count={}".format(width * height * 4)
            ])
            logging.info(f"裁剪截图已成功保存到设备路径: {output_path}")
        except Exception as e:
            logging.error(f"裁剪截图失败: {e}")
            raise

    def pull_screenshot(self, device_path, local_path):
        """
        将截图文件从设备拉取到本地。
        :param device_path: 设备上的截图路径
        :param local_path: 本地保存路径
        :raises: 如果拉取失败，抛出异常并记录日志
        """
        logging.info(f"开始将截图从设备路径 {device_path} 拉取到本地路径 {local_path}")
        try:
            self.run_adb_command(["pull", device_path, local_path])
            logging.info(f"截图已成功保存到本地路径: {local_path}")
        except Exception as e:
            logging.error(f"拉取截图失败: {e}")
            raise

    def delete_file(self, device_path):
        """
        删除设备上的文件。
        :param device_path: 设备上的文件路径
        :raises: 如果删除失败，抛出异常并记录日志
        """
        logging.info(f"开始删除设备上的文件: {device_path}")
        try:
            self.run_adb_command(["shell", "rm", device_path])
            logging.info(f"设备上的文件已成功删除: {device_path}")
        except Exception as e:
            logging.error(f"删除文件失败: {e}")
            raise

    def take_partial_screenshot(self, x, y, width, height, local_output_path):
        """
        捕获部分截图并保存到本地。
        :param x: 裁剪区域的左上角X坐标
        :param y: 裁剪区域的左上角Y坐标
        :param width: 裁剪区域的宽度
        :param height: 裁剪区域的高度
        :param local_output_path: 本地保存路径
        :raises: 如果操作失败，抛出异常并记录日志
        """
        logging.info(f"开始捕获部分截图，裁剪区域: ({x}, {y}, {width}, {height})，保存到本地路径: {local_output_path}")
        try:
            # 1. 捕获全屏截图
            full_screenshot_path = "/sdcard/fullscreen.png"
            self.take_full_screenshot(full_screenshot_path)

            # 2. 裁剪截图
            cropped_screenshot_path = "/sdcard/part.png"
            self.crop_screenshot(full_screenshot_path, cropped_screenshot_path, x, y, width, height)

            # 3. 将裁剪后的截图拉取到本地
            self.pull_screenshot(cropped_screenshot_path, local_output_path)

            # 4. 清理设备上的临时文件
            self.delete_file(full_screenshot_path)
            self.delete_file(cropped_screenshot_path)

            logging.info(f"部分截图已成功保存到本地路径: {local_output_path}")
        except Exception as e:
            logging.error(f"捕获部分截图失败: {e}")
            raise